import numpy as np
import DRW_library as dl
import emcee
import celerite
from celerite import terms
from scipy.optimize import minimize
from multiprocessing import Pool
from time import perf_counter

start=perf_counter()
    
def mainprocessing(m):
    tau=2922/1.047**m
    re=np.array([])
    for k in range(N):       
        np.random.seed(seed[m*N+k])
        #t=np.linspace(0,2922,epochs)
        t=np.sort(np.random.randint(0,2922,epochs)+
                  np.random.uniform(-0.13,0.13,epochs))
        y=dl.DRW_process(t,tau,0.2,18)
        ls_sdss=[]
        lsigma=[]
        for i in range(epochs):
            s=np.sqrt(0.004**2+np.exp(1.63*(y[i]-22.55)))
            #s=np.sqrt(0.013**2+np.exp(2*(y[i]-23.36)))
            lsigma.append(s)
            ls_sdss.append(np.random.normal(y[i],s,1)[0])
        s_sdss=np.array(ls_sdss)
        sigma=np.array(lsigma)
        re=np.append(re,dl.DRW_fit(t,s_sdss,sigma,method=method))
    re.resize((N,2))  
    print('[',m,']',np.log10(2922/tau))
    return np.array([2922/tau,np.mean(re[:,0])/tau,np.std(re[:,0])/tau,
                     np.mean(re[:,1])/0.1414,np.std(re[:,1])/0.1414])
    
#BEGIN
epochs=60
N=1000
method='MAX'
seed=np.random.randint(1,2**31-1,200*N)
print('epochs=',epochs,'N=',N,'method=',method)
if __name__ == '__main__':
    pool=Pool(6)
    l=np.array(pool.map(mainprocessing, np.arange(0,150,4)))
    pool.close()
    pool.join()    
np.savetxt('f60.txt',l,fmt='%.4f')
print('time=',perf_counter()-start)
#END    
